import joblib
import typing
# 多线程优化模块 (有线程池)
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed
from itertools import *
from base_funcs import *


params: dict = joblib.load("params.pkl")
max_epoch = params["T"] - 1
alpha = params["alpha"]
K_t = params["K_t"]
p_K_t = params["p_K_t"]
D_t = params["D_t"]
p_D_t = params["p_D_t"]
p = params["p"]
b = params["b"]
c_e = params["c_e"]
c_m = params["c_m"]
h_e = params["h_e"]
h_m = params["h_m"]
theta = params["theta"]

_max_worker = 5
_thread_accelerate = True
_process_accelerate = False

print("base_model_funcs loaded")


def grid_search_exa(iter_grid, func: typing.Callable, params: tuple, epoch: int) -> list:
    if _thread_accelerate:
        if _process_accelerate:
            pool = ProcessPoolExecutor(max_workers=_max_worker)
        else:
            pool = ThreadPoolExecutor(max_workers=_max_worker)
        tasks = [pool.submit(func, *params, x, y, epoch) for x, y in iter_grid]
        results = [t.result() for t in as_completed(tasks)]
        return results
    else:
        return [func(*params, x, y, epoch) for x, y in iter_grid]


def linear_search_exa(iter_arr, func: typing.Callable, params: tuple, epoch: int) -> list:
    if _thread_accelerate:
        if _process_accelerate:
            pool = ProcessPoolExecutor(max_workers=_max_worker)
        else:
            pool = ThreadPoolExecutor(max_workers=_max_worker)
        tasks = [pool.submit(func, *params, x, epoch) for x in iter_arr]
        results = [t.result() for t in as_completed(tasks)]
        return results
    else:
        return [func(*params, x, y, epoch) for x, y in iter_arr]


# 基础算子
def psi_1T(v_1, v_2, v_3, e, z):
    return h_e * E(I(v_1 - e - D_t), p_D_t) \
           + b * E(I(D_t - v_1 + e), p_D_t) \
           + c_e * (v_1 - e) \
           + c_m * (z - v_3) \
           + h_m * (v_3 - v_1) \
           + theta * e


def psi_2T(v_2, v_3, y, z):
    return h_e * E(I(y - D_t), p_D_t) \
           + b * E(I(D_t - y), p_D_t) \
           + c_e * y \
           + c_m * (z - v_3) \
           + h_m * (v_3 - y)


def psi_3T(v_3, y, z):
    return h_e * E(I(y - D_t), p_D_t) \
           + b * E(I(D_t - y), p_D_t) \
           + c_e * y \
           + c_m * (z - v_3) \
           + h_m * (v_3 - y)


# psi_1t
def psi_1t(v_1, v_2, v_3, e, z, epoch):
    # 先计算基础值
    val = psi_1T(v_1, v_2, v_3, e, z)
    if epoch != max_epoch:
        # (-2, 3, 4, 5)
        _D_t = [f_t(v_1 - e - D_t_i, v_2 - e - D_t_i, v_3 - e - D_t_i, z - e - D_t_i, epoch + 1) for D_t_i in D_t]
        val += alpha * E(_D_t, p_D_t)
    return val


# psi_2t
def psi_2t(v_2, v_3, y, z, epoch):
    val = psi_2T(v_2, v_3, y, z)
    if epoch != max_epoch:
        _D_t = [f_t(y - D_t_i, v_2 - D_t_i, v_3 - D_t_i, z - D_t_i, epoch + 1) for D_t_i in D_t]
        val += alpha * E(_D_t, p_D_t)
    return val


# psi_3t
def psi_3t(v_3, y, z, epoch):
    val = psi_3T(v_3, y, z)
    if epoch != max_epoch:
        _D_t = [f_t(y - D_t_i, y - D_t_i, v_3 - D_t_i, z - D_t_i, epoch + 1) for D_t_i in D_t]
        val += alpha * E(_D_t, p_D_t)
    return val


def phi_1t(v_0, v_1, v_2, v_3, K_t_i, epoch):
    e_range = range(0, v_1 - v_0 + 1) # [1, 2]
    z_range = range(v_3, v_3 + K_t_i + 1)  # [3, 4]
    e_z_grid = product(e_range, z_range) # [(1, 3), (1, 4), (2, 3), (2, 4)]
    psi_1t_values = grid_search_exa(e_z_grid, psi_1t, (v_1, v_2, v_3), epoch)
    return - c_e * v_0 + np.min(psi_1t_values)


def phi_2t(v_0, v_1, v_2, v_3, K_t_i, epoch):
    y_range = range(v_1, v_2 + 1)
    z_range = range(v_3, v_3 + K_t_i + 1)
    y_z_grid = product(y_range, z_range)
    psi_2t_values = grid_search_exa(y_z_grid, psi_2t, (v_2, v_3), epoch)
    return - c_e * v_0 + np.min(psi_2t_values)


def phi_3t(v_0, v_1, v_2, v_3, K_t_i, epoch):
    y_range = range(v_2, v_3 + 1)
    z_range = range(v_3, v_3 + K_t_i + 1)
    y_z_grid = product(y_range, z_range)
    psi_3t_values = grid_search_exa(y_z_grid, psi_3t, (v_3,), epoch)
    return - c_e * v_0 + np.min(psi_3t_values)


def f_t(v_0, v_1, v_2, v_3, epoch):
    _K_t = []
    for K_t_i in K_t:
        exp_1_value = phi_1t(v_0, v_1, v_2, v_3, K_t_i, epoch)
        exp_2_value = phi_2t(v_0, v_1, v_2, v_3, K_t_i, epoch)
        exp_3_value = phi_3t(v_0, v_1, v_2, v_3, K_t_i, epoch)
        _K_t.append(np.min([exp_1_value, exp_2_value, exp_3_value]))
    return E(_K_t, p_K_t)
